from pyblog import *
import numpy

## model
@var_dist
def cluster_wts(): return InfDirichlet(ALPHA)

@var_dist
def cluster(i): return InfCategorical(cluster_wts())

@var_dist
def mean(k): return MNormal(numpy.matrix([0.0, 0.0]).transpose(),
                            numpy.matrix([[1e-6, 0], [0, 1e-6]]))

# here precision is the inverse covariance
@var_dist
def prec(k): return Wishart(numpy.matrix([[20000.0, 10000.0],
                                          [10000.0, 20000.0]]), 4)

@var_dist
def obs_data(i): return MNormal(mean(cluster(i)), prec(cluster(i)))


### model specific helper functions ###

def all_mean_prec_wts(world):
  """
  Returns the means and precisions and weights of all the clusters used
  in the world
  """
  return [(world[mean(k)], world[prec(k)], world[cluster_wts()][k]) \
          for k in set(value for rvar, value in world.iteritems()
                       if rvar.fn_name() == cluster)]

def generate(alpha, numpoints):
  """
  Returns a list of mean and precision pairs followed by all the
  generated points
  """
  # set the hyper-parameters
  global ALPHA
  ALPHA = alpha
  
  ans = query([obs_data(i) for i in range(numpoints)] + [pyblog_world()], [],
              burnin=0, scans=0, stats=False, outtyp = QUERY_LAST)

  assert(len(ans) == 1 + numpoints)
  clusters = all_mean_prec_wts(ans[-1])
  train = ans[:numpoints/2]
  test = ans[numpoints/2:-1]

  return train, test, clusters

def estimate(train, scans, last_world=None):
  """
  Returns a list of means and precisions
  """
  inits = None
  if last_world is not None:
    inits = [rvar == value for rvar, value in last_world.iteritems()]

  statsobj = {}
  world, = query([pyblog_world()],
                 [obs_data(i)==v for i,v in enumerate(train)], inits,
                 burnin = 0, scans=scans, stats=False,
                 outtyp = QUERY_LAST, statsobj=statsobj)
  
  return all_mean_prec_wts(world), world, statsobj['scan-time']
